import numpy as np
from tensorflow import keras

num_classes = 10
input_shape = (28, 28, 1)

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# LeNet-5
model = keras.Sequential()
model.add(keras.layers.Conv2D(filters=6, kernel_size=(5, 5), input_shape=(28, 28, 1), padding='same',
                                 activation="sigmoid"))
model.add(keras.layers.AveragePooling2D(pool_size=(2, 2)))
model.add(keras.layers.Conv2D(filters=16, kernel_size=(5, 5), activation="sigmoid"))
model.add(keras.layers.AveragePooling2D(pool_size=(2, 2)))
model.add(keras.layers.Conv2D(filters=120, kernel_size=(5, 5), activation="sigmoid"))
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(84, activation='sigmoid'))
model.add(keras.layers.Dense(10, activation='softmax'))
model.summary()

# 设置优化器、损失函数、准确率
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['acc'])
# 载入训练集、测试级，设置训练轮次
history = model.fit(x_train, y_train, batch_size=128, epochs=30, validation_data=(x_test, y_test))
# 使用测试集进行评估
model.evaluate(x_test, y_test)
# 保存模型
model.save('mnist.h5')
